package dht

import (
	"errors"
	"log"
	"sync"
	"time"

	"gopkg.in/mgo.v2/bson"
)

type Provide struct {
	Hash  ID                      `bson:"hash"`
	Peers map[string]*ProvidePeer `bson:"peers"`
}

type ProvidePeer struct {
	Addr     string `bson:"addr"`
	ExpireAt int64  `bson:"expire_at"`
}

type ProvideReq struct {
	Hash ID     `bson:"hash"`
	Addr string `bson:"addr"`
}

func (dht *DHT) provideHandler(reqb []byte) ([]byte, error) {
	req := &ProvideReq{}
	err := bson.Unmarshal(reqb, &req)
	if err != nil {
		return nil, err
	}
	provide := &Provide{}
	data, ok := dht.store.Get("provide:" + req.Hash.String())
	if ok == true {
		err := bson.Unmarshal([]byte(data), &provide)
		if err != nil {
			return nil, err
		}
	} else {
		provide.Hash = req.Hash
		provide.Peers = map[string]*ProvidePeer{}
	}
	provide.Peers[req.Addr] = &ProvidePeer{
		Addr:     req.Addr,
		ExpireAt: time.Now().Add(time.Second * 3).Unix(),
	}
	provideb, err := bson.Marshal(provide)
	if err != nil {
		return nil, err
	}
	err = dht.store.Put("provide:"+provide.Hash.String(), string(provideb))
	if err != nil {
		return nil, err
	}
	return nil, nil
}

func (dht *DHT) provide(hash ID, reqAddr string, provideAddr string) error {
	req := &ProvideReq{
		Hash: hash,
		Addr: provideAddr,
	}
	reqb, err := bson.Marshal(req)
	if err != nil {
		return err
	}
	_, err = dht.query(reqAddr, PROVIDE, reqb)
	if err != nil {
		return err
	}
	return nil
}

func (dht *DHT) Provide(hash ID, reqAddr string, provideAddr string) error {
	list, err := dht.FindNode(reqAddr, hash)
	if err != nil {
		return err
	}
	if len(list) == 0 {
		return errors.New("There is no enough node")
	}
	var wg sync.WaitGroup
	for _, node := range list {
		wg.Add(1)
		go func(node *Node) {
			defer wg.Done()
			err = dht.provide(hash, node.Addr, provideAddr)
			if err != nil {
				log.Println("ERROR:", err)
				return
			}
		}(node)
	}
	wg.Wait()
	return nil
}
